import os
import xml.etree.ElementTree as ET

def convert_voc_to_yolo(xml_folder, output_folder):
    # 获取xml文件列表
    xml_files = [f for f in os.listdir(xml_folder) if f.endswith('.xml')]

    for xml_file in xml_files:
        xml_path = os.path.join(xml_folder, xml_file)
        tree = ET.parse(xml_path)
        root = tree.getroot()

        image_name = root.find('filename').text
        image_width = int(root.find('size/width').text)
        image_height = int(root.find('size/height').text)

        txt_file = os.path.splitext(xml_file)[0] + '.txt'
        txt_path = os.path.join(output_folder, txt_file)

        with open(txt_path, 'w') as f:
            for obj in root.findall('object'):
                class_name = obj.find('name').text
                class_id = 0  # 如果有多个类别，这里需要根据具体情况进行修改
                bbox = obj.find('bndbox')
                x_min = int(bbox.find('xmin').text)
                y_min = int(bbox.find('ymin').text)
                x_max = int(bbox.find('xmax').text)
                y_max = int(bbox.find('ymax').text)

                x_center = (x_min + x_max) / 2 / image_width
                y_center = (y_min + y_max) / 2 / image_height
                width = (x_max - x_min) / image_width
                height = (y_max - y_min) / image_height

                f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

# 用法示例
xml_folder = 'C:/Users/admin/Desktop/SafetyHelmetWearingDataset/Safety Helmet Wearing Dataset/9rcv8mm682-1/Safety Helmet Wearing Dataset/Annotations'
output_folder = 'C:/Users/admin/Desktop/SafetyHelmetWearingDataset/Safety Helmet Wearing Dataset/9rcv8mm682-1/Safety Helmet Wearing Dataset/labels'
convert_voc_to_yolo(xml_folder, output_folder)
